-
Notifications
You must be signed in to change notification settings - Fork 139
Document inplace_on_inputs method with working example and remove TODO comment #1575
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Conversation
Op.inplace_on_inputs
Co-authored-by: ricardoV94 <[email protected]>
Op.inplace_on_inputs
and returns a new `Op` instance that will perform inplace operations only on | ||
those inputs where it is safe and beneficial to do so. | ||
|
||
.. testcode:: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Instead of this code snippet provide a minimal self-contained example with a simple Op that is Blockwise, and show that if we set the input to be Mutable the inplace version will be introduced. Something like:
import pytensor
from pytensor.graph.op import Op
from pytensor.tensor.blockwise import Blockwise
class MyOpWithInplace(Op):
__props__ = ("destroy_a",)
def __init__(self, destroy_a):
self.destroy_a = destroy_a
if destroy_a:
self.destroy_map = {0: [0]}
def make_node(self, a):
return Apply(self, [a], [a.type()]
def perform(self, node, inputs, output_storage):
[a] = inputs
if not self.destroy_a:
a = a.copy()
a[0] += 1
output_storage[0][0] = a
# Fill this in
def inplace_on_inputs(...):
...
a = pt.vector("a")
# Only Blockwise trigger inplace automatically for now
# Since the Blockwise isn't needed in this case, it will be removed after the inplace optimization
op = Blockwise(MyOpWithInplace(destroy_a=False), signature="(a)->=(a)")
out = op(a)
# Give PyTensor permission to inplace on user provided inputs
fn = pytensor.function([pytensor.In(a, mutable=True)], out)
# Confirm that we have the inplace version of the Op
fn.dprint(print_destroy_map=True)
test_a = np.zeros(5)
fn(test_a)
print(test_a) # Should be modified in place
Then create the expected test output as well, which should include both the dprint, and test_a
with a 1 after the first call. Break the code so the dprint check comes before the array check.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I've replaced the simple code snippet with a complete self-contained example as requested. The new example includes:
- A full
MyOpWithInplace
class implementinginplace_on_inputs
properly - Demonstration using Blockwise to trigger automatic inplace optimization
- Expected dprint output showing the inplace version (
destroy_a=True
) is used - Test showing actual inplace array modification behavior
The code is broken into two testcode blocks as requested - first showing the dprint output, then the array modification test. Both include corresponding testoutput blocks with expected results.
Commit: a0f8eff
…on_inputs Co-authored-by: ricardoV94 <[email protected]>
This PR addresses the documentation of the
inplace_on_inputs
method as requested in the issue. The method allowsOp
classes to create inplace versions of themselves for specified inputs, enabling memory optimizations while maintaining correctness.Changes Made
pytensor/graph/op.py
in theinplace_on_inputs
method definitiondoc/extending/inplace.rst
explaining theinplace_on_inputs
methodDocumentation Added
The new documentation section includes:
MyOpWithInplace
class showing:inplace_on_inputs
implementationThe example demonstrates how PyTensor automatically calls
inplace_on_inputs
when optimizing Blockwise operations, and shows the resulting memory-efficient inplace behavior when inputs are marked as mutable.This method is automatically used by PyTensor's optimization pipeline and is crucial for efficient memory management in Blockwise operations. The documentation makes it clear that users typically don't need to call this method directly, but Op developers should understand how to implement it properly.
💬 Share your feedback on Copilot coding agent for the chance to win a $200 gift card! Click here to start the survey.